'''
Function:
    Define the l2norm
Author:
    Zhenchao Jin
'''
import numpy as np
import luojianet
from luojianet import nn, ops, Parameter, Tensor
from luojianet.common.initializer import initializer, Constant
from luojianet.common import initializer as init
import luojianet.ops as ops

'''L2Norm'''
class L2Norm(nn.Module):
    def __init__(self, channels, scale=10, eps=1e-10):
        super(L2Norm, self).__init__()
        self.channels, self.eps = channels, eps
        # self.weight = nn.Parameter(torch.Tensor(channels))
        # self.weight = Parameter(Tensor(np.random.randn(channels), dtype=luojianet.float32), init=Tensor(scale, dtype=luojianet.float32))
        self.weight = Parameter(initializer(Constant(scale), channels, luojianet.float32))
        # nn.init.constant_(self.weight, scale)
    '''forward'''
    def forward(self, x):
        # norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps
        norm = x.pow(2).sum(axis=1, keepdims=True).sqrt() + self.eps
        # x = torch.div(x, norm)
        x = ops.div(x, norm)
        out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(x) * x
        return out